/*
 * Treebank.java
 *
 * Created on August 10, 2006, 5:40 PM
 *
 */

package galronnlp.perceptron;

import galronnlp.io.PennTreeReader;
import galronnlp.util.ProductionGenerator;
import galronnlp.util.Symbol;
import galronnlp.util.SymbolGenerator;
import galronnlp.perceptron.Grammar;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.util.Iterator;
import java.util.LinkedList;

/**
 *
 * @author Daniel A. Galron
 */
public class Treebank extends LinkedList<Node> {
    
    public static final int FULL     = 0;
    public static final int SAMPLE   = 1;
    public static final int WSJ      = 2;
    public static final int BROWN    = 3;

    
    protected int sample = FULL;
    protected String CorpusRoot = "";
    protected int corpus = WSJ;
    
    protected SymbolGenerator sg;
    protected int from, to;
    
    
    /** Creates a new instance of Treebank */
    public Treebank(String root, int sample, int corpus, SymbolGenerator sg, int fromSection, int toSection) {
        if(root.charAt(root.length() - 1) != '/') {
            CorpusRoot = root + "/";
        } else {
            CorpusRoot = root;
        }
        this.sample = sample;
        this.corpus = corpus;
        this.sg = sg;        
        this.from = fromSection;
        this.to = toSection;
        if(this.sample==this.SAMPLE) {
            this.from = 1;
            this.to = 1;
        }
    }
    /** Creates a new instance of Treebank */
    public Treebank(String root, int sample, int corpus, SymbolGenerator sg) {
        this(root, sample, corpus, sg, 0, 24);
    }
    
    public Treebank(String root, SymbolGenerator sg) {
        this(root, FULL, WSJ, sg);
    }

    public void build() {
        this.build(this.from, this.to);
    }
    
    public void build(int from, int to) {
        if(corpus == WSJ)
            this.readInWSJFromTo(from, to);
        else if(corpus == BROWN)
            System.err.println("NOT YET IMPLEMENTED");
    }
    
    private void readInWSJ() {
        if(sample == FULL) this.readInWSJFull();
        else if(sample == SAMPLE) this.readInWSJSample();
    }
    
    private void readInWSJFull() {
        readInWSJFromTo(2,21);
    }
    
    private void readInWSJFromTo(int from, int to) {
        int fileNum = 1;
        int dirNum  = 0;
        CorpusRoot += "/mrg/wsj/";
        String fileExt = ".mrg.transformed";
        for(dirNum = from; dirNum < (to+1); dirNum++) {
            String dirString = "";
            if(dirNum < 10) dirString = "0" + Integer.toString(dirNum);
            else dirString = Integer.toString(dirNum);
            if(dirNum == 0) {
                for(fileNum = 1; fileNum < 100; fileNum++) {
                    String fileString = "";
                    if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
                    else fileString = Integer.toString(fileNum);
                    System.err.println(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                    this.readInWSJTreeFile(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                }
            }
            else if(dirNum == 3) {
                for(fileNum = 0; fileNum < 81; fileNum++) {
                    String fileString = "";
                    if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
                    else fileString = Integer.toString(fileNum);
                    System.err.println(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                    this.readInWSJTreeFile(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                }
                
            }
            else if(dirNum == 8) {
                for(fileNum = 0; fileNum < 21; fileNum++) {
                    String fileString = "";
                    if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
                    else fileString = Integer.toString(fileNum);
                    System.err.println(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                    this.readInWSJTreeFile(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                }
                
            }
            else if(dirNum == 21) {
                for(fileNum = 0; fileNum < 73; fileNum++) {
                    String fileString = "";
                    if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
                    else fileString = Integer.toString(fileNum);
                    System.err.println(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                    this.readInWSJTreeFile(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                }
            }
            else if(dirNum == 24) {
                for(fileNum = 0; fileNum < 55; fileNum++) {
                    String fileString = "";
                    if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
                    else fileString = Integer.toString(fileNum);
                    System.err.println(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                    this.readInWSJTreeFile(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                }
            }
            else if(dirNum == 22) {
                for(fileNum = 0; fileNum < 83; fileNum++) {
                    String fileString = "";
                    if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
                    else fileString = Integer.toString(fileNum);
                    System.err.println(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                    this.readInWSJTreeFile(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                }
            }
            else {
                for(fileNum = 0; fileNum < 100; fileNum++) {
                    String fileString = "";
                    if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
                    else fileString = Integer.toString(fileNum);
                    System.err.println(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                    this.readInWSJTreeFile(CorpusRoot + dirString + "/wsj_" + dirString + fileString + fileExt);
                }
            }
        }
    }
    
    private void readInWSJSample() {
        CorpusRoot += "/mrg/wsj/";
        for(int fileNum = 1; fileNum < 100; fileNum++) {
            String fileString = "";
            if(fileNum < 10) fileString = "0" + Integer.toString(fileNum);
            else fileString = Integer.toString(fileNum);
//            System.err.println("wsj_00" + fileString + ".mrg");
            this.readInWSJTreeFile(CorpusRoot + "01" + "/wsj_" + "01" + fileString + ".mrg.transfomed");
        }
    }
    
    private void readInWSJTreeFile(String file) {
        try {
            PennTreeReader reader = new PennTreeReader(new BufferedReader(new FileReader(file)), sg);
            while(reader.hasNext()) {
                Node tree = reader.readNodeTree();
                if(tree != null) {
                    //System.err.println(tree.toTreeString());
                    this.addLast(tree);
                }
            }
        } catch(Exception e) {
            e.printStackTrace();
        }
    }

    public String toString() {
        String ret = "---------------------------------------\n";
        for(Iterator<Node> i = this.iterator(); i.hasNext(); ) {
            ret += i.next();
            ret += "---------------------------------------\n";
        }
        return ret;
    }
    
    public static void main(String[] args) {
        SymbolGenerator sg = new SymbolGenerator();
        Treebank bank = new Treebank(args[0], FULL, WSJ, sg, Integer.parseInt(args[1]), Integer.parseInt(args[2]));
        bank.build();
        System.out.println("Done! Treebank size = " + bank.size());
        Grammar g = new Grammar(new ProductionGenerator(), sg);
        Iterator<Node> it = bank.iterator();
        Node n = null;
        while(it.hasNext()) {
            n = it.next();
            if(n!=null)
                n.addToGrammar(g);
        }
        g.writeXML(new File(args[0]+"../grammar01.gram"));
        System.out.println("---------------------------------------\n");        
        System.exit(0);
    }
}
